package com.example.asyncquery.service;

import com.example.asyncquery.model.QueryResult;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
@Slf4j
public class SseEmitterService {
    
    // 存储SSE连接的缓存，以queryId为键
    private final Map<String, SseEmitter> emitters = new ConcurrentHashMap<>();
    
    /**
     * 为特定查询创建一个新的SSE连接
     * @param queryId 查询ID
     * @return SseEmitter实例
     */
    public SseEmitter createEmitter(String queryId) {
        // 创建一个长时间存活(30分钟)的发射器
        SseEmitter emitter = new SseEmitter(1800000L);
        
        // 设置超时回调
        emitter.onTimeout(() -> {
            log.info("SSE连接超时: {}", queryId);
            emitter.complete();
            emitters.remove(queryId);
        });
        
        // 设置完成回调
        emitter.onCompletion(() -> {
            log.info("SSE连接完成: {}", queryId);
            emitters.remove(queryId);
        });
        
        // 设置错误回调
        emitter.onError(ex -> {
            log.error("SSE连接错误: {}, 原因: {}", queryId, ex.getMessage());
            emitter.complete();
            emitters.remove(queryId);
        });
        
        // 保存到缓存
        emitters.put(queryId, emitter);
        log.info("创建新的SSE连接: {}", queryId);
        
        return emitter;
    }
    
    /**
     * 发送查询结果更新
     * @param queryId 查询ID
     * @param result 查询结果
     */
    public void sendQueryResultUpdate(String queryId, QueryResult<?> result) {
        SseEmitter emitter = emitters.get(queryId);
        if (emitter != null) {
            try {
                log.info("发送查询结果更新: {}, 状态: {}", queryId, result.getStatus());
                emitter.send(SseEmitter.event()
                        .id(queryId)
                        .name("query-update")
                        .data(result));
                
                // 如果查询已完成或失败，则完成SSE连接
                if (result.getStatus().toString().equals("COMPLETED") || 
                    result.getStatus().toString().equals("FAILED")) {
                    log.info("查询已{}，关闭SSE连接: {}", result.getStatus(), queryId);
                    emitter.complete();
                    emitters.remove(queryId);
                }
            } catch (IOException e) {
                log.error("发送SSE事件失败: {}", e.getMessage());
                emitter.completeWithError(e);
                emitters.remove(queryId);
            }
        } else {
            log.warn("未找到查询的SSE连接: {}", queryId);
        }
    }
    
    /**
     * 移除SSE连接
     * @param queryId 查询ID
     */
    public void removeEmitter(String queryId) {
        SseEmitter emitter = emitters.remove(queryId);
        if (emitter != null) {
            emitter.complete();
            log.info("已移除SSE连接: {}", queryId);
        }
    }
} 